package com.demo.security.utils;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.UUID;

public class RSAUtils {

    final static String ALGORITHM = "RSA";
    private static final int DEFAULT_KEY_SIZE = 2048;

    /**
     * 从文件中读取公钥
     *
     * @param filename 公钥保存路径，相对于classpath
     * @return 公钥对象
     * @throws IOException
     * @throws InvalidKeySpecException
     */
    public static PublicKey getPublicKey(String filename) throws IOException, InvalidKeySpecException {
        byte[] bytes = readFile(filename);
        return getPublicKey(bytes);
    }

    /**
     * 从文件中读取密钥
     *
     * @param filename 私钥保存路径，相对于classpath
     * @return 私钥对象
     * @throws IOException
     * @throws InvalidKeySpecException
     */
    public static PrivateKey getPrivateKey(String filename) throws IOException, InvalidKeySpecException {
        byte[] bytes = readFile(filename);
        return getPrivateKey(bytes);
    }

    /**
     * 根据密文，生存rsa公钥和私钥,并写入指定文件
     *
     * @param publicKeyFilename  公钥文件路径
     * @param privateKeyFilename 私钥文件路径
     * @param secret             生成密钥的密文
     * @throws NoSuchAlgorithmException
     * @throws IOException
     */
    public static void generateKey(String publicKeyFilename, String privateKeyFilename, String secret, int keySize) throws IOException {
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
            SecureRandom secureRandom = new SecureRandom(secret.getBytes());
            keyPairGenerator.initialize(Math.max(keySize, DEFAULT_KEY_SIZE), secureRandom);
            KeyPair keyPair = keyPairGenerator.genKeyPair();
            // 获取公钥并写出
            byte[] publicKeyBytes = keyPair.getPublic().getEncoded();
            publicKeyBytes = Base64.getEncoder().encode(publicKeyBytes);
            writeFile(publicKeyFilename, publicKeyBytes);
            System.out.println(">>> PublicKey: " + publicKeyFilename);
            // 获取私钥并写出
            byte[] privateKeyBytes = keyPair.getPrivate().getEncoded();
            privateKeyBytes = Base64.getEncoder().encode(privateKeyBytes);
            writeFile(privateKeyFilename, privateKeyBytes);
            System.out.println(">>> PrivateKey: " + privateKeyFilename);
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 获取公钥
     *
     * @param bytes 公钥的字节形式
     * @return
     * @throws InvalidKeySpecException
     */
    private static PublicKey getPublicKey(byte[] bytes) throws InvalidKeySpecException {
        KeyFactory factory;
        try {
            factory = KeyFactory.getInstance(ALGORITHM);
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
        bytes = Base64.getDecoder().decode(bytes);
        X509EncodedKeySpec spec = new X509EncodedKeySpec(bytes);
        return factory.generatePublic(spec);
    }

    /**
     * 获取密钥
     *
     * @param bytes 私钥的字节形式
     * @return
     * @throws InvalidKeySpecException
     */
    private static PrivateKey getPrivateKey(byte[] bytes) throws InvalidKeySpecException {
        KeyFactory factory;
        try {
            factory = KeyFactory.getInstance(ALGORITHM);
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
        bytes = Base64.getDecoder().decode(bytes);
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(bytes);
        return factory.generatePrivate(spec);
    }

    private static byte[] readFile(String fileName) throws IOException {
        byte[] bytes = null;
        InputStream is = RSAUtils.class.getResourceAsStream(fileName);
        if (is == null) {
            bytes = Files.readAllBytes(new File(fileName).toPath());
        } else {
            bytes = toByteArray(is);
        }
        return bytes;
    }

    private static void writeFile(String destPath, byte[] bytes) throws IOException {
        File dest = new File(destPath);
        if (!dest.exists()) {
            dest.createNewFile();
        }
        Files.write(dest.toPath(), bytes);
    }

    private static byte[] toByteArray(InputStream input) throws IOException {
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        byte[] buffer = new byte[4096];
        int n = 0;
        while (-1 != (n = input.read(buffer))) {
            output.write(buffer, 0, n);
        }
        return output.toByteArray();
    }

    public static void main(String[] args) throws IOException, InvalidKeySpecException {
        long now = System.currentTimeMillis();
        String secret = UUID.randomUUID().toString();
        String localPath = "D:\\user\\01418295\\桌面\\tmp\\";
        String publicKeyFilename = localPath + "RSA_" + now + ".pub";
        String privateKeyFilename = localPath + "RSA_" + now + ".key";
        generateKey(publicKeyFilename, privateKeyFilename, secret, DEFAULT_KEY_SIZE);

        PublicKey publicKey = getPublicKey("/cert/RSA_1637893836223.pub");
        System.out.println(publicKey.getFormat());
        PrivateKey privateKey = getPrivateKey(localPath + "RSA_1637893836223.key");
        System.out.println(privateKey.getFormat());
    }
}